package ka

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/Shopify/sarama"
	"github.com/gogf/gf/container/gvar"
	"github.com/gogf/gf/encoding/gbase64"
	"go.opentelemetry.io/contrib"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/label"
	"go.opentelemetry.io/otel/semconv"
	"go.opentelemetry.io/otel/trace"
	"math/rand"
	"strconv"
	"sync"
)

type ConsumerConfig struct {
	// broker的集群地址
	Brokers []string `toml:"brokers"`
	// topic 的名称
	Topic string `toml:"topic"`
	// 消费组名称
	Group string `toml:"group"`
	SASL  struct {
		Enable   bool   `toml:"enable"`
		User     string `toml:"user"`
		Password string `toml:"password"`
	} `toml:"sasl"`
	// 多少个协程
	Workers int `toml:"workers"`
	// 是否从最老的开始消费
	Oldest bool `toml:"oldest"`
}

func (c *ConsumerConfig) Check() bool {
	if len(c.Topic) == 0 ||
		len(c.Brokers) == 0 {
		return false
	}
	if c.Group == "" {
		c.Group = fmt.Sprintf("%s%.8x", gTestPrefix, rand.Int63())
	}
	if c.Workers == 0 {
		c.Workers = 10
	}
	return true
}

func NewConsumer(cfg *ConsumerConfig) (*Consumer, error) {
	return newConsumerWithConfig(cfg)
}

func newConsumerWithConfig(cfg *ConsumerConfig) (*Consumer, error) {
	if !cfg.Check() {
		return nil, errors.New("config error")
	}
	saramaCfg := sarama.NewConfig()
	{
		saramaCfg.Consumer.Return.Errors = true
		saramaCfg.Version = GetVersion()
		if cfg.Oldest {
			saramaCfg.Consumer.Offsets.Initial = sarama.OffsetOldest
		} else {
			saramaCfg.Consumer.Offsets.Initial = sarama.OffsetNewest
		}
		if cfg.SASL.Enable {
			saramaCfg.Net.SASL.Enable = cfg.SASL.Enable
			saramaCfg.Net.SASL.User = cfg.SASL.User
			saramaCfg.Net.SASL.Password = cfg.SASL.Password
		}
	}
	if err := saramaCfg.Validate(); err != nil {
		return nil, err
	}
	// Start with a client
	client, err := sarama.NewClient(cfg.Brokers, saramaCfg)
	if err != nil {
		return nil, err
	}
	// Start a new consumer group
	consumerGroup, err := sarama.NewConsumerGroupFromClient(cfg.Group, client)
	if err != nil {
		return nil, err
	}
	c := new(Consumer)
	c.cfg = cfg
	c.saramaCfg = saramaCfg
	c.group = consumerGroup
	c.wg = &sync.WaitGroup{}
	c.ctx, c.cancel = context.WithCancel(context.Background())
	c.consumerEventHandlers = make(map[string]HandleConsumerFunc)
	c.handlerLock = sync.RWMutex{}

	c.tracer = otel.GetTracerProvider().Tracer(gConsumerTracerName,
		trace.WithInstrumentationVersion(contrib.SemVersion()))

	return c, nil
}

type consumerGroupHandler struct {
	consumer *Consumer
}

func (h consumerGroupHandler) getTopic() string {
	return h.consumer.GetTopic()
}

func (consumerGroupHandler) Setup(_ sarama.ConsumerGroupSession) error   { return nil }
func (consumerGroupHandler) Cleanup(_ sarama.ConsumerGroupSession) error { return nil }
func (h consumerGroupHandler) ConsumeClaim(sess sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
	for {
		select {
		case msg, ok := <-claim.Messages():
			if !ok {
				return nil
			}

			e := &Event{}
			if err := json.Unmarshal(msg.Value, e); err != nil {
				// 这里解析错误
				sess.MarkMessage(msg, gMsgDone)
				if h.consumer.e != nil {
					h.consumer.e(err)
				}
				continue
			}
			parentSpanContext := otel.GetTextMapPropagator().Extract(context.Background(), e)
			// Create a span.
			attrs := []label.KeyValue{
				semconv.MessagingSystemKey.String("kafka"),
				semconv.MessagingDestinationKindKeyTopic,
				semconv.MessagingDestinationKey.String(msg.Topic),
				semconv.MessagingOperationReceive,
				semconv.MessagingMessageIDKey.String(strconv.FormatInt(msg.Offset, 10)),
				gKafkaPartitionKey.Int32(msg.Partition),
			}
			opts := []trace.SpanOption{
				trace.WithAttributes(attrs...),
				trace.WithSpanKind(trace.SpanKindConsumer),
			}

			newCtx, span := h.consumer.tracer.Start(parentSpanContext, "kafka.consume", opts...)
			h.consumer.handlerLock.RLock()
			handler, ok := h.consumer.consumerEventHandlers[e.Type]
			h.consumer.handlerLock.RUnlock()

			b, _ := gbase64.DecodeString(e.Data.String())
			e.Data = gvar.New(b)

			if ok {
				if err := handler(newCtx, e); err != nil {
					return err
				}
			}
			if h.consumer.consumerTopicHandler != nil {
				_ = h.consumer.consumerTopicHandler(newCtx, e)
			}

			span.End()
			sess.MarkMessage(msg, gMsgDone)
		}
	}
}

func (c *Consumer) HandleError(e HandleErrorFunc) {
	c.e = e
}

func (c *Consumer) GetTopic() string {
	return c.cfg.Topic
}

func (c *Consumer) HandleEvent(eventType string, consumerFunc HandleConsumerFunc) {
	c.handlerLock.Lock()
	defer c.handlerLock.Unlock()
	c.consumerEventHandlers[eventType] = consumerFunc
}

func (c *Consumer) HandleTopic(handler HandleConsumerFunc) {
	c.handlerLock.Lock()
	defer c.handlerLock.Unlock()
	c.consumerTopicHandler = handler
}

func (c *Consumer) Run() {
	c.handlerLock.Lock()
	defer c.handlerLock.Unlock()
	if c.run {
		return
	}
	go c.Handle()
	for i := 0; i < c.cfg.Workers; i++ {
		go c.Worker()
	}
	c.run = true
}

func (c *Consumer) Close() error {
	c.cancel()
	c.wg.Wait()
	return c.group.Close()
}

func (c *Consumer) Worker() {
	c.wg.Add(1)
	defer c.wg.Done()
	topics := []string{c.cfg.Topic}
	handler := consumerGroupHandler{c}
	for {
		if err := c.group.Consume(c.ctx, topics, handler); err != nil {
			if c.e != nil {
				c.e(err)
			}
		}

		if c.ctx.Err() != nil {
			return
		}
	}
}

func (c *Consumer) Handle() {
	for {
		select {
		case err, ok := <-c.group.Errors():
			{
				if !ok {
					return
				}
				if c.e != nil {
					c.e(err)
				}
			}
		}
	}
}
